# -*- coding: utf-8 -*-
'''
Created on 07.11.2019

@author: yu03
'''

import numpy as np
import os
import re
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import glob
from mpl_toolkits.mplot3d import Axes3D
from Video_Unified import folder_path, np_result_names, hor_index, ver_index, hor_lines, ver_lines
import sys
from FFT_Interpolation import FFT_interpolation_nonlinearity_compare
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from collections import OrderedDict

plt_show_mode = 'show' ### plt.show()
# plt_show_mode = 'save' ### plt.save()

i = 0
fs_cam = 50

''' 
    Nonlinearity FFT Compare
''' 
fig1 = plt.figure('Nonlinearity FFT Compare horxver')
plt.gcf().set_size_inches(18,9)
ax1 = fig1.add_subplot(2, 2, 1)
ax2 = fig1.add_subplot(2, 2, 2)
ax3 = fig1.add_subplot(2, 2, 3)
ax4 = fig1.add_subplot(2, 2, 4)

''' 
    for all groups
''' 
for np_result in np_result_names:
# for np_result in np_result_names[3:6]:
    ''' 
        reading .npy file
    ''' 
    file_name = np_result.split('\\')[-1] ### x_lines.py
#     file_name = re.findall(r"openmode\\(.+).npy", np_result)[0]
    print(file_name)
    i += 1
    f = open(np_result, 'rb')
    f_size = os.fstat(f.fileno()).st_size
    ''' 
        put 4-D results in "lines"
    ''' 
    lines = []
    time_sequence = []
    while f.tell() < f_size:
        line = np.load(f, allow_pickle=True)
        lines.append(line)
    print('%i: %s, %eB'%(i, file_name, f_size))
    lines = np.array(lines)[::,::,300:2200]
    print(np.shape(lines), 'Channle, lines, frames:')
    frame_num = np.shape(lines)[2]
    line_num = np.shape(lines)[1]
    
    hor_angle_set, ver_angle_set, hor_length_set, ver_length_set = lines
    ''' 
        making average
    ''' 
    averaged_results = []
    
    hor_angle_avr = np.mean(hor_angle_set, axis=0)
    ver_angle_avr = np.mean(ver_angle_set, axis=0)
    hor_length_avr = np.mean(hor_length_set, axis=0)
    ver_length_avr = np.mean(ver_length_set, axis=0)
    
    averaged_results.append(hor_angle_avr)
    averaged_results.append(ver_angle_avr)
    averaged_results.append(hor_length_avr)
    averaged_results.append(ver_length_avr)
        
    ''' 
        polyfit degree
    '''     
    length_dof = 10
    angle_dof = 6

    ''' 
        fitting averaged results
    '''
    averaged_nonlinearity = []
    for data in averaged_results[0:2]:
        data_fit_params = np.polyfit(np.arange(frame_num), data, angle_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        nonlinearity = data - data_fitline
        averaged_nonlinearity.append(nonlinearity)
    for data in averaged_results[2:4]:
        data_fit_params = np.polyfit(np.arange(frame_num), data, length_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        nonlinearity = data - data_fitline
        averaged_nonlinearity.append(nonlinearity)

    print(np.shape(averaged_nonlinearity))
    [hor_angle_avr_nonlinearity, ver_angle_avr_nonlinearity, hor_length_avr_nonlinearity, ver_length_avr_nonlinearity] = averaged_nonlinearity
    print(np.shape(hor_angle_avr_nonlinearity))

    ''' 
        fitting linebyline results
    '''
    hor_angle_nonlinearity_set = []
    ver_angle_nonlinearity_set = []
    hor_length_nonlinearity_set = []
    ver_length_nonlinearity_set = []    
      
    for i in range(line_num):
#         for i in [50,51,52]:
        data = hor_angle_set[i]
        data_fit_params = np.polyfit(np.arange(frame_num), data, angle_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        hor_angle_nonlinearity = data - data_fitline
           
           
        data = hor_length_set[i]
        data_fit_params = np.polyfit(np.arange(frame_num), data, length_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        hor_length_nonlinearity = data - data_fitline
           
        data = ver_angle_set[i]
        data_fit_params = np.polyfit(np.arange(frame_num), data, angle_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        ver_angle_nonlinearity = data - data_fitline
           
        data = ver_length_set[i]
        data_fit_params = np.polyfit(np.arange(frame_num), data, length_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        ver_length_nonlinearity = data - data_fitline
           
        hor_angle_nonlinearity_set.append(hor_angle_nonlinearity)
        hor_length_nonlinearity_set.append(hor_length_nonlinearity)
        ver_angle_nonlinearity_set.append(ver_angle_nonlinearity)
        ver_length_nonlinearity_set.append(ver_length_nonlinearity)
           
           
    hor_angle_nonlinearity_set = np.array(hor_angle_nonlinearity_set)
    hor_length_nonlinearity_set = np.array(hor_length_nonlinearity_set)
    ver_angle_nonlinearity_set = np.array(ver_angle_nonlinearity_set)
    ver_length_nonlinearity_set = np.array(ver_length_nonlinearity_set)
       
       
    hor_angle_nonlinearity_max = np.amax(hor_angle_nonlinearity_set, axis=1)
    hor_angle_nonlinearity_min = np.amin(hor_angle_nonlinearity_set, axis=1)
    hor_length_nonlinearity_max = np.amax(hor_length_nonlinearity_set, axis=1)
    hor_length_nonlinearity_min = np.amin(hor_length_nonlinearity_set, axis=1)
       
    ver_angle_nonlinearity_max = np.amax(ver_angle_nonlinearity_set, axis=1)
    ver_angle_nonlinearity_min = np.amin(ver_angle_nonlinearity_set, axis=1)
    ver_length_nonlinearity_max = np.amax(ver_length_nonlinearity_set, axis=1)
    ver_length_nonlinearity_min = np.amin(ver_length_nonlinearity_set, axis=1)
    
    
    hor_non_fft_set = np.empty(shape=[101, 150])
    ver_non_fft_set = np.empty(shape=[101, 150])
    hor_length_non_fft_set = np.empty(shape=[101, 150])
    ver_length_non_fft_set = np.empty(shape=[101, 150])
    
    hor_non_phi_set = np.empty(shape=[101, 150])
    ver_non_phi_set = np.empty(shape=[101, 150])
    hor_length_non_phi_set = np.empty(shape=[101, 150])
    ver_length_non_phi_set = np.empty(shape=[101, 150])
    
    for i in np.arange(101):
        hor_non_f, hor_non_phi, freqline, hor_non_fft = FFT_interpolation_nonlinearity_compare(hor_angle_nonlinearity_set[int(i), 200:1000], 1/fs_cam)
        hor_non_fft_set[i] = np.array(hor_non_fft[20:170])
        hor_non_phi_set[i] = np.array(hor_non_phi[20:170])
        ver_non_f, ver_non_phi, freqline, ver_non_fft = FFT_interpolation_nonlinearity_compare(ver_angle_nonlinearity_set[int(i), 200:1000], 1/fs_cam)
        ver_non_fft_set[i] = np.array(ver_non_fft[20:170])
        hor_length_non_f, hor_length_non_phi, freqline, hor_length_non_fft = FFT_interpolation_nonlinearity_compare(hor_length_nonlinearity_set[int(i), 200:1000], 1/fs_cam)
        hor_length_non_fft_set[i] = np.array(hor_length_non_fft[20:170])
        ver_length_non_f, ver_length_non_phi, freqline, ver_length_non_fft = FFT_interpolation_nonlinearity_compare(ver_length_nonlinearity_set[int(i), 200:1000], 1/fs_cam)
        ver_length_non_fft_set[i] = np.array(ver_length_non_fft[20:170])
    
    hor_non_fft_set = np.array(hor_non_fft_set)
    ver_non_fft_set = np.array(ver_non_fft_set)
    hor_length_non_fft_set = np.array(hor_length_non_fft_set)
    ver_length_non_fft_set = np.array(ver_length_non_fft_set)
    
    hor_non_phi_set = np.array(hor_non_phi_set)
    ver_non_phi_set = np.array(ver_non_phi_set)
    hor_length_non_phi_set = np.array(hor_length_non_phi_set)
    ver_length_non_phi_set = np.array(ver_length_non_phi_set)
    
    f1_set, f2_set = [], []
    for k in range(100):
        f1 = np.where(hor_non_fft_set[k][20:40]==np.max(hor_non_fft_set[k][20:40]))[0][0] + 20
        f2 = np.where(hor_non_fft_set[k][70:90]==np.max(hor_non_fft_set[k][70:90]))[0][0] + 70
        f1_set.append(f1)
        f2_set.append(f2)
    f1 = max(f1_set,key=f1_set.count)
    f2 = max(f2_set,key=f2_set.count)
    f_1 = np.min([f1,f2])
    f_2 = np.max([f1,f2])
    print(f_1, f_2)

    hor_non_phi_1, hor_non_phi_2 = hor_non_phi_set[:, f_1], hor_non_phi_set[:, f_2]
        
    ax1.plot(hor_lines, hor_non_fft_set[:, f_1], marker='o', color='blue', markersize='2', label='f1')
    ax1.plot(hor_lines, hor_non_fft_set[:, f_2], marker='o', color='red', markersize='2', label='f2')
#     ax1.plot(hor_lines, (hor_non_phi_set[:, f_1]), marker='o', color='black', markersize='2', label='phi1')
#     ax1.plot(hor_lines, (hor_non_phi_set[:, f_2]), marker='o', color='green', markersize='2', label='phi2')
    
    ax2.plot(ver_lines, ver_non_fft_set[:, f_1], marker='o', color='blue', markersize='2', label='f1')
    ax2.plot(ver_lines, ver_non_fft_set[:, f_2], marker='o', color='red', markersize='2', label='f2')
#     ax2.plot(ver_lines, (ver_non_phi_set[:, f_1]), marker='o', color='black', markersize='2', label='phi1')
#     ax2.plot(ver_lines, (ver_non_phi_set[:, f_2]), marker='o', color='green', markersize='2', label='phi2')
    
    ax3.plot(hor_lines, hor_length_non_fft_set[:, f_1], color='blue', marker='o', markersize='2', label='f1')
    ax3.plot(hor_lines, hor_length_non_fft_set[:, f_2], color='red', marker='o', markersize='2', label='f2')
#     ax3.plot(hor_lines, (hor_length_non_phi_set[:, f_1]), marker='o', color='black', markersize='2', label='phi1')
#     ax3.plot(hor_lines, (hor_length_non_phi_set[:, f_2]), marker='o', color='green', markersize='2', label='phi2')
    
    ax4.plot(ver_lines, ver_length_non_fft_set[:, f_1], color='blue', marker='o', markersize='2', label='f1')
    ax4.plot(ver_lines, ver_length_non_fft_set[:, f_2], color='red', marker='o', markersize='2', label='f2')
#     ax4.plot(ver_lines, (ver_length_non_phi_set[:, f_1]), marker='o', color='black', markersize='2', label='phi1')
#     ax4.plot(ver_lines, (ver_length_non_phi_set[:, f_2]), marker='o', color='green', markersize='2', label='phi2')

ax1.title.set_text('Hor. Tilt')
ax2.title.set_text('Ver. Tilt')
ax3.title.set_text('Hor. Displacement')
ax4.title.set_text('Ver. Displacement')

ax1.set_xlabel('Hor. Line Num.')
ax2.set_xlabel('Ver. Line Num.')
ax3.set_xlabel('Hor. Line Num.')
ax4.set_xlabel('Ver. Line Num.')

ax1.set_ylabel('Amplitude / urad')
ax2.set_ylabel('Amplitude / urad')
ax3.set_ylabel('Amplitude / nm')
ax4.set_ylabel('Amplitude / nm')

ax1.grid(which='both', axis='both')
ax2.grid(which='both', axis='both')
ax3.grid(which='both', axis='both')
ax4.grid(which='both', axis='both')

handles, labels = plt.gca().get_legend_handles_labels()
by_label = OrderedDict(zip(labels, handles))
ax1.legend(by_label.values(), by_label.keys(), loc='upper right')
ax2.legend(by_label.values(), by_label.keys(), loc='upper right')
ax3.legend(by_label.values(), by_label.keys(), loc='upper right')
ax4.legend(by_label.values(), by_label.keys(), loc='upper right')

plt.tight_layout()
if plt_show_mode == 'save':
    plt.savefig(folder_path + '\\' + 'Nonlinearity_FFT_Compare_90x8.png', dpi=300)
    plt.close()
elif plt_show_mode == 'show':
    plt.show()
    pass
else:
    sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)   

''' 
    Nonlinearity FFT
''' 
fig = plt.figure('Nonlinearity FFT')
plt.gcf().set_size_inches(18,9)
ax1 = fig.add_subplot(2, 2, 1, projection='3d')
ax2 = fig.add_subplot(2, 2, 2, projection='3d')
ax3 = fig.add_subplot(2, 2, 3, projection='3d')
ax4 = fig.add_subplot(2, 2, 4, projection='3d')
plt.gca().patch.set_facecolor('white')
X_fft = np.arange(101)
#     Y_fft = freqline[20:170]
Y_fft = np.arange(150)
X_fft, Y_fft = np.meshgrid(Y_fft, X_fft)
surf1 = ax1.plot_surface(X_fft, Y_fft, hor_non_fft_set, cmap='jet', rstride=1, cstride=1,
                       linewidth=0, antialiased=False)
surf2 = ax2.plot_surface(X_fft, Y_fft, ver_non_fft_set, cmap='jet', rstride=1, cstride=1,
                       linewidth=0, antialiased=False)
surf3 = ax3.plot_surface(X_fft, Y_fft, hor_length_non_fft_set, cmap='jet', rstride=1, cstride=1,
                       linewidth=0, antialiased=False)
surf4 = ax4.plot_surface(X_fft, Y_fft, ver_length_non_fft_set, cmap='jet', rstride=1, cstride=1,
                       linewidth=0, antialiased=False)
 
ax1.yaxis.set_major_formatter(FormatStrFormatter('%i'))
ax2.yaxis.set_major_formatter(FormatStrFormatter('%i'))
ax3.yaxis.set_major_formatter(FormatStrFormatter('%i'))
ax4.yaxis.set_major_formatter(FormatStrFormatter('%i'))
 
ax1.title.set_text('Horizontal Tilt')
ax2.title.set_text('Vertical Tilt')
ax3.title.set_text('Horizontal Length')
ax4.title.set_text('Vertical Length')
    
ax1.set_ylabel('Hor. Line Num.')
ax1.set_xlabel('Freq. / Hz')
ax1.set_zlabel('Amplitude / urad')
    
ax2.set_ylabel('Ver. Line Num.')
ax2.set_xlabel('Freq. / Hz')
ax2.set_zlabel('Amplitude / urad')
 
ax3.set_ylabel('Hor. Line Num.')
ax3.set_xlabel('Freq. / Hz')
ax3.set_zlabel('Amplitude / nm')
    
ax4.set_ylabel('Ver. Line Num.')
ax4.set_xlabel('Freq. / Hz')
ax4.set_zlabel('Amplitude / nm')
    
ax1.view_init(60, -60)
ax2.view_init(60, -60)
ax3.view_init(60, -60)
ax4.view_init(60, -60)

plt.tight_layout()
if plt_show_mode == 'save':
#     plt.savefig(folder_path + '\\' + 'Nonlinearity_FFT.png', dpi=300)
    plt.close()
elif plt_show_mode == 'show':
    plt.show()
    pass
else:
    sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)   